# -*- coding: utf-8 -*-

__author__ = "Kamil Gołuński"
__date__ = "$2009-11-08 22:54:54$"



from copy import deepcopy
from numpy import *
import os
from time import clock
import wx
from wx.lib.floatcanvas import FloatCanvas
from wx.lib.floatcanvas import NavCanvas
from wx.lib.floatcanvas.Utilities import *

from classes import *
import svq
ID_DRAW_BUTTON = 100
ID_QUIT_BUTTON = 101
ID_CLEAR_BUTTON = 103
ID_ZOOM_IN_BUTTON = 104
ID_ZOOM_OUT_BUTTON = 105
ID_ZOOM_TO_FIT_BUTTON = 110
ID_MOVE_MODE_BUTTON = 111
ID_TEST_BUTTON = 112

ID_ABOUT_MENU = 200
ID_EXIT_MENU  = 201
ID_ZOOM_IN_MENU = 202
ID_ZOOM_OUT_MENU = 203
ID_ZOOM_TO_FIT_MENU = 204
ID_DRAWTEST_MENU = 205
ID_DRAWMAP_MENU = 206
ID_CLEAR_MENU = 207
ID_OPEN = 208
ID_DIR = 209
ID_RUNCBSVQ_MENU = 210
ID_FINAL_MATCH = 211

ID_TEST = 500


class DrawFrame(wx.Frame):
    def __init__(self, parent, id, title, position, size):
        wx.Frame.__init__(self, parent, id, title, position, size)

        ## Set up the MenuBar

        MenuBar = wx.MenuBar()

        file_menu = wx.Menu()
        file_menu.Append(ID_OPEN, "O&pen", "Open bitmap")
        file_menu.Append(ID_DIR, "Open directory\tCtrl+O")
        file_menu.Append(ID_EXIT_MENU, "E&xit", "Terminate the program")
        wx.EVT_MENU(self, ID_EXIT_MENU, self.OnQuit)
        wx.EVT_MENU(self, ID_OPEN, self.OnOpen)
        wx.EVT_MENU(self, ID_DIR, self.OnOpenDir)
        MenuBar.Append(file_menu, "&File")

        draw_menu = wx.Menu()

        draw_menu.Append(ID_DRAWTEST_MENU, "Show code-book points\tF5", "")
        wx.EVT_MENU(self, ID_DRAWTEST_MENU, self.updateCanvas)
        draw_menu.Append(ID_DRAWMAP_MENU, "Run SVQ on each image\tF6", "")
        wx.EVT_MENU(self, ID_DRAWMAP_MENU, self.onStart)
        draw_menu.Append(ID_RUNCBSVQ_MENU, "Run SVQ on codebook vectors\tF7", "")
        wx.EVT_MENU(self, ID_RUNCBSVQ_MENU, self.onSecondStart)
        draw_menu.Append(ID_FINAL_MATCH, "Final matching\tF8", "")
        wx.EVT_MENU(self, ID_FINAL_MATCH, self.FinalMatching)
        #wx.EVT_MENU(self, ID_FINAL_MATCH, self.test)
        draw_menu.Append(ID_CLEAR_MENU, "&Clear", "Clear the canvas")
        wx.EVT_MENU(self, ID_CLEAR_MENU, self.Clear)
        MenuBar.Append(draw_menu, "&Run")


        help_menu = wx.Menu()
        help_menu.Append(ID_ABOUT_MENU, "&About",
                         "More information About this program")
        wx.EVT_MENU(self, ID_ABOUT_MENU, self.OnAbout)
        MenuBar.Append(help_menu, "&Help")

        self.SetMenuBar(MenuBar)

        self.CreateStatusBar()
        self.SetStatusText("")

        wx.EVT_CLOSE(self, self.OnCloseWindow)

        wx.EVT_RIGHT_DOWN(self, self.RightButtonEvent)

        #self.Canvas = NavCanvas.NavCanvas(self, -1, (750, 600),
        #                                  Debug=False,
        #                                  BackgroundColor="Black").Canvas
        self.Canvas = FloatCanvas.FloatCanvas(self, -1, (750,600),Debug=False,
                                          BackgroundColor="Black")
        self.Canvas.NumBetweenBlits = 1000
        self.Show(True)

        #self.DrawSVQ(None)

        ## Algorithm parameters
        panel = wx.Panel(self, -1)
        self.textTimes = wx.TextCtrl(panel,-1,"",size=(50, 23),pos = (120,10))
        self.textCodeBook = wx.TextCtrl(panel,-1,"",size=(50, 23),pos=(120,40))
        self.textBeta = wx.TextCtrl(panel,-1,"",size=(50, 23),pos=(280,10))
        self.textEta = wx.TextCtrl(panel,-1,"",size=(50, 23),pos= (280,40))
        #self.textFrom = wx.TextCtrl(panel,-1,"",size=(50, 23),pos=(100,30))
        #self.textTo = wx.TextCtrl(panel,-1,"",size=(50, 23),pos=(250,30))

        self.textTimes.SetValue('5000')
        self.textBeta.SetValue('0.002')
        self.textEta.SetValue('0.5')
        self.textCodeBook.SetValue('20')
        #self.textFrom.SetValue('200')
        #self.textTo.SetValue('400')

        self.labelTimes = wx.StaticText(panel,-1,"Iteracje ",pos=(6,15))
        self.labelCodeBook = wx.StaticText(panel,-1,"Punkty kotwicowe",pos=(6,45))
        self.labelBeta = wx.StaticText(panel,-1,"Beta ", pos = (236,20))
        self.labelEta = wx.StaticText(panel,-1,"Eta ", pos = (236,45))
        #self.labelFrom = wx.StaticText(panel,-1,"P. kotwicowe od",pos=(6,40))
        #self.labelTo = wx.StaticText(panel,-1,"P. kotwicowe do",pos=(154,40))

        self.start = wx.Button(panel,888,"Start",pos=(500,10))
        self.stop = wx.Button(panel,999,"Stop",pos=(500,40))

        self.Bind(wx.EVT_BUTTON, self.onStart, id=888)
        self.Bind(wx.EVT_BUTTON, self.onStop, id=999)

        box = wx.BoxSizer(wx.VERTICAL)
        hbox = wx.BoxSizer(wx.HORIZONTAL)
        hbox.Add(panel,0,0,0)

        box.Add(hbox,0,0,0)

        box.Add(self.Canvas, 1, wx.EXPAND,0 )


        self.SetSizer(box)
        self.Layout()
        self.stop_cond = False
        self.prefix= 'cyfryA-'
        return None

    def onStart(self,event):
        self.times = int(self.textTimes.GetValue())
        self.beta = float(self.textBeta.GetValue())
        self.eta = float(self.textEta.GetValue())
        self.c = int(self.textCodeBook.GetValue())
        self.RunSVQ(event)
        #import time
        #start = time.time()
        #self.doSVQ(t,b,e,c,o,d)
        #print time.time() - start

    def onSecondStart(self,event):
        self.times = int(self.textTimes.GetValue())
        self.beta = float(self.textBeta.GetValue())
        self.eta = float(self.textEta.GetValue())
        self.c = int(self.textCodeBook.GetValue())
        self.RunCbSVQ(event)

    def onStop(self,event):
        print 'stop'
        self.stop_cond = True


    def RightButtonEvent(self, event):
        print "Right Button has been clicked in DrawFrame"
        print "coords are: %i, %i" % (event.GetX(), event.GetY())
        event.Skip()

    def OnAbout(self, event):
        dlg = wx.MessageDialog(self, "-----------About---------", "About Me", wx.OK | wx.ICON_INFORMATION)
        dlg.ShowModal()
        dlg.Destroy()

    def ZoomToFit(self, event):
        self.Canvas.ZoomToBB()

    def Clear(self, event=None):
        self.Canvas.ClearAll()
        self.Canvas.Draw()

    def OnQuit(self, event):
        self.Close(True)

    def OnCloseWindow(self, event):
        self.Destroy()

    def OnOpen(self, event):
        dialog = wx.FileDialog (None, style=wx.OPEN)
        if dialog.ShowModal() == wx.ID_OK:
            self.bmp = wx.Bitmap(dialog.GetPath())
        dialog.Destroy()
        self.InitSVQ()

    def OnOpenDir(self, event):
        dialog = wx.DirDialog(None, defaultPath='C:\Users\Kamil\Documents\NetBeansProjects\kamillo-ma-thesis-svq\series')
        if dialog.ShowModal() == wx.ID_OK:
            self.images = self.ReadAllImages(dialog.GetPath())
            self.AddImagesToCanvas(self.images)
        self.cbvecs = [self.Canvas.AddPointSet([[-10000, -10000]], Color='Red', Diameter=3, InForeground=False) for i in range(len(self.images))]
        
        dialog.Destroy()

    def AddImagesToCanvas(self, images):
        position = [-375, 300]
        #750 / 5
        for img in images:
            self.Canvas.AddScaledBitmap(img.image, position[:], 300, 'tl', InForeground=False)
            position[0] += img.image.GetWidth()
            if (images.index(img) + 1) % 5 == 0:
                position[0] = -375
                position[1] -= img.image.GetHeight()
        self.Canvas.ZoomToBB()

    def MovePoints(self, points, position):
        new_points = []
        for i in range(len(points)):
            new_points.append([points[i][0] + position[0], points[i][1] * (-1) + position[1]])
        return new_points

    def updateCanvas(self, canvas_points, bb=None, event=None):
        position = [-375, 300]
        Canvas = self.Canvas
        for cb in range(len(canvas_points)):
            self.cbvecs[cb].SetPoints(self.MovePoints(canvas_points[cb], position[:]))
            position[0] += self.images[cb].image.GetWidth()
            if (cb + 1) % 5 == 0:
                position[0] = -375
                position[1] -= self.images[cb].image.GetHeight()
        Canvas.ZoomToBB(bb)

    def GenerateColors(self,n):
        colors = [(random.randrange(256),random.randrange(256),random.randrange(256)) for i in range(n)]
        return colors

    def CalcBB(self,points):
        position = [-375, 300]
        new_points = []
        for cb in range(len(points)):
            new_points.append(self.MovePoints(points[cb], position[:]))
            position[0] += self.images[cb].image.GetWidth()
            if (cb + 1) % 4 == 0:
                position = [-375, 300]
        bb = BBox.fromPoints(new_points)
        return bb


    def RunSVQ(self, event=None):
        times = self.times
        beta = self.beta
        eta = self.eta
        eta0 = eta
        c = self.c
        start = clock()
        #position = [-375, 300]
        images = self.images
        cbvectors = [CodeBookVector(s, c, 'P') for s in images]
        self.updateCanvas([cbvectors[i].getPoints() for i in range(len(cbvectors))])
        for state in svq.svq(images, times, eta, beta, cbvectors, 'shape'):
            wx.GetApp().Yield(True)
            self.updateCanvas([cbvectors[i].getPoints() for i in range(len(cbvectors))])
            if self.stop_cond == True:
                break
        self.Canvas.SaveAsImage(self.prefix+str(eta0)+'-'+str(beta)+'first.png')
        print "running the svq algorithm took %f seconds" % (clock() - start)
        #print cbvectors[0].getPoints()
        self.cbvectors = cbvectors
        if self.stop_cond == False:
            self.Bind(wx.EVT_BUTTON, self.onSecondStart, id=888)
        self.stop_cond = False
        

    def RunCbSVQ(self, event=None):
        times = self.times
        beta = self.beta
        eta = self.eta
        c = self.c
        self.Canvas.ClearAll()
        colors = self.GenerateColors(c)
        #colors = ["Red", "Blue", "Green", "Yellow"]
        points = [cb.cbvector for cb in self.cbvectors]
        self.shapevec = Shape().fromVector(points)
        cbv = CodeBookVector(self.shapevec, c, 'S')
        self.cbvecs = [self.Canvas.AddPointSet(cbv.getPoints()[i], Color=colors[i], Diameter=3, InForeground=False) for i in range(len(cbv.getPoints()))]
        self.updateCanvas(cbv.getPoints())
        bb = self.CalcBB([self.cbvectors[i].getPoints() for i in range(len(self.cbvectors))])
        
        bb=self.Canvas.ZoomToBB()
        #for state in svq.svq([self.shapevec],times,0.5,0.00002,[cbv],'iteration'):
        #    wx.GetApp().Yield(True)
        #    self.updateCanvas(cbv.getPoints())
        eta0 = eta
        start = clock()
        self.Canvas.SaveAsImage(self.prefix+str(eta0)+'-'+str(beta)+'.png')
        self.images[0].image.GetWidth()
        for i in range(times):  # iterations of algorithm
            if i%100 == 0:
                self.Canvas.SaveAsImage(self.prefix+str(eta0)+'-'+str(beta)+'-'+str(i)+'.png')
            point = self.shapevec.pickRandomPoint()
            
            cbv.munkresMinMatching(point)
            f = cbv.calcF(point, beta, True)
            cbv.move(point, f, eta)
            eta = eta0 * (times - i) / times
            self.updateCanvas(cbv.getPoints(),bb)
            wx.GetApp().Yield(True)
            if self.stop_cond == True:
                break
        print clock() - start
        self.updateCanvas(cbv.getPoints(),bb)
        self.stop_cond = False
        self.cbv = cbv

    def test(self,event=None):
        cbv = CodeBookVector(self.shapevec, len(self.shapevec.shape_points), 'S')
        cbv.initFinalMatchingMatrix(len(self.shapevec.shape_points))

    def FinalMatching(self,event=None):
        cbv = self.cbv
        shapevec = self.shapevec
        cbvectors = self.cbvectors

        cbv.initFinalMatchingMatrix(len(shapevec.shape_points))
        for shape in shapevec.shape_points:
            cbv.munkresMinMatching(shape,True)
        m = Munkres()
        ind = m.compute(cbv.final_matching_matrix)
        indx = []
        for i in range(len(cbv.final_matching_matrix[0])):
            tmp = [cbv.final_matching_matrix[j][i] for j in range(len(cbv.final_matching_matrix))]
            a = min(tmp)
            indx.append(tmp.index(a))
        #for i in indx:
        #    self.cbvecs.append(self.Canvas.AddPointSet(self.cbvectors[i].getPoints(), Color="Red", Diameter=3, InForeground=True) )
        #t = self.Canvas.AddPointSet(cbv.getPoints(), Color = "White", Diameter=3, InForeground=True )
        s = [self.Canvas.AddPointSet(self.cbvectors[i].getPoints(), Color="Red", Diameter=3, InForeground=False) for i in indx]
        z = [self.cbvectors[i].getPoints() for i in indx]
        #self.updateCanvas([self.cbvectors[i].getPoints() for i in indx])
        
        position = [-375, 300]
        Canvas = self.Canvas
        for cb in range(len(z)):
            s[cb].SetPoints(self.MovePoints(z[cb], position[:]))
            position[0] += self.images[cb].image.GetWidth()
            if (cb + 1) % len(z) == 0:
                position[0] = -375
                position[1] -= self.images[cb].image.GetHeight()
        for i in range(len(self.cbvecs)):
            self.cbvecs[i].SetColor(wx.Color(255,255,255,150))
        Canvas.ZoomToBB()
        self.Canvas.SaveAsImage(self.prefix+str(self.eta)+'-'+str(self.beta)+'final.png')

    def ReadAllImages(self, dir):
        images = []
        for subdir, dir, files in os.walk(dir):
            progress_dlg = wx.ProgressDialog(title='Open', message='Loading: ' + str(files[0]), maximum=len(files), style=wx.PD_AUTO_HIDE | wx.PD_APP_MODAL | wx.PD_SMOOTH)
            files.sort()
            for file in files:
                if '.svn' in dir:
                    dir.remove('.svn')
                images.append(Shape().fromPath(os.path.join(subdir, file)))
                progress_dlg.Update(files.index(file) + 1, 'Loading: ' + str(file))
                wx.GetApp().Yield(True)
        progress_dlg.Destroy()
        return images

class DemoApp(wx.App):
    def OnInit(self):
        wx.InitAllImageHandlers()
        frame = DrawFrame(None, -1, "Soft Vector Quantization", wx.DefaultPosition, (750, 700))

        self.SetTopWindow(frame)

        return True


if __name__ == "__main__":
    try:
        import psyco
        psyco.log()
        psyco.full()
    except:
        print "could not import psyco"

    app = DemoApp(0)
    app.MainLoop()











